package drds.plus.datasource.api;

import drds.plus.datasource.connection_restrict.ConnectionRestrictHelper;
import drds.plus.datasource.connection_restrict.ConnectionRestrictSlot;
import lombok.extern.slf4j.Slf4j;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;

@Slf4j
public class ConnectionWrapper {
    private final String applicationId;
    private final DataSourceWrapper dataSourceWrapper;
    private final Connection connection;
    private ConnectionRestrictSlot connectionRestrictSlot;
    private Set<StatementWrapper> statementWrapperSet = Collections.synchronizedSet(new HashSet<StatementWrapper>(1));


    public ConnectionWrapper(String applicationId, DataSourceWrapper dataSourceWrapper, Connection connection, ConnectionRestrictSlot connectionRestrictSlot) {
        this.applicationId = applicationId;
        this.dataSourceWrapper = dataSourceWrapper;
        this.connection = connection;
        this.connectionRestrictSlot = connectionRestrictSlot;


    }

    public Connection getConnection() {
        return connection;
    }

    public void close() throws SQLException {
        try {
            this.connection.close();
        } finally {
            if (connectionRestrictSlot != null) {
                connectionRestrictSlot.release();
                connectionRestrictSlot = null; // 防止重复关闭
            }
            dataSourceWrapper.threadCount.decrementAndGet();
            ConnectionRestrictHelper.removeConnRestrictKey();
        }

        for (StatementWrapper statementWrapper : this.statementWrapperSet) {
            try {
                statementWrapper.close(false);
            } catch (SQLException e) {
                log.error("", e);
            }
        }
    }

    void removeStatementWrapper(StatementWrapper statementWrapper) {
        if (!statementWrapperSet.remove(statementWrapper)) {
            log.warn("current statmenet ：" + statementWrapper + " doesn't exist!");
        }
    }

    public void commit() throws SQLException {
        this.connection.commit();
    }

    public StatementWrapper createStatement() throws SQLException {
        Statement statement = this.connection.createStatement();
        StatementWrapper statementWrapper = new StatementWrapper(statement, this, this.dataSourceWrapper, this.applicationId);
        this.statementWrapperSet.add(statementWrapper);
        return statementWrapper;
    }

    public boolean getAutoCommit() throws SQLException {
        return this.connection.getAutoCommit();
    }

    public void setAutoCommit(boolean autoCommit) throws SQLException {
        this.connection.setAutoCommit(autoCommit);
    }

    public String getCatalog() throws SQLException {
        return this.connection.getCatalog();
    }

    public void setCatalog(String catalog) throws SQLException {
        this.connection.setCatalog(catalog);
    }

    public boolean isClosed() throws SQLException {
        return this.connection.isClosed();
    }

    public PreparedStatementWrapper prepareStatement(String sql) throws SQLException {
        PreparedStatement preparedStatement = this.connection.prepareStatement(sql);
        PreparedStatementWrapper preparedStatementWrapper = new PreparedStatementWrapper(preparedStatement, this, this.dataSourceWrapper, sql, this.applicationId);
        statementWrapperSet.add(preparedStatementWrapper);
        return preparedStatementWrapper;
    }

    public PreparedStatementWrapper prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
        PreparedStatement preparedStatement = this.connection.prepareStatement(sql, autoGeneratedKeys);
        PreparedStatementWrapper preparedStatementWrapper = new PreparedStatementWrapper(preparedStatement, this, this.dataSourceWrapper, sql, this.applicationId);
        statementWrapperSet.add(preparedStatementWrapper);
        return preparedStatementWrapper;
    }

    public void rollback() throws SQLException {
        this.connection.rollback();
    }


}
